﻿using FaissOperation.FaissIndex;
using FaissOperation.FaissIndex.Index;
using FaissOperation.Interface;
using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Text.Json;
using System.Threading.Tasks;

namespace FaissOperation.FaissIndexFlat
{
    public class IVFVectorCollection : BaseFaissCollection
    {
        private const int DEFAULT_NLIST = 128; // 默认桶数量


        public override IEnumerable<SearchResult> Search(float[] query, int topK)
        {
            // 维度不匹配，不会返回任何数据。
            if (query.Length != Dimension)
            {
                return [];
            }

            var q = NormalizeOnAddAndQuery
                ? FaissHelper.NormalizeForFaiss(query)
                : query;

            var (labels, distances) = faissIndex.Search(q, topK);
            var results = new List<SearchResult>(vectors.Count);
            for (int i = 0; i < labels.Length; i++)
            {
                if (vectors.ContainsKey((int)labels[i]))
                    results.Add(new SearchResult((int)labels[i], distances[i]));
            }
            return results.ToArray();
        }

        public override void CreateIndex()
        {
            var quantizer = new FlatIndex(Dimension, FaissHelper.ToFaissMetric(MetricTypeInUse));

            faissIndex = new IVFIndex(quantizer, Dimension, DEFAULT_NLIST, FaissHelper.ToFaissMetric(MetricTypeInUse));
        }

    }


}
